# coding : UTF-8
"""
作者：BingBO   时间：2022年10月23日
自动调整代码格式 ：Alt+Ctrl+L
"""
import cv2 as cv
import numpy as np
import math
from module_V import fitLine
from numpy.linalg import solve  # 解线性方程


def laser1_Extract(L1img, start_point1, EP1):
    print("********************** laser1_ 特征提取 **********************")
    Red = L1img[:, :, 2]  # B0 G1 R2
    _, threshold = cv.threshold(Red, 190, 255, cv.THRESH_BINARY)  # 转为二值图
    laser1_threshold = cv.medianBlur(threshold, 3)

    # cv.imshow('laser1_threshold', laser1_threshold)
    cv.imwrite(r"F:\MVS_Data\laser1_threshold.jpg", laser1_threshold)

    # 求laser1光条的左间断点列坐标
    lightList = np.sum(laser1_threshold, axis=0)  # 0 列相加, 变成一行
    ret = np.where(lightList > 0)  # 返回（索引+ 数据类型 ）
    lightIdx = ret[0]
    # print("lightIdx:", lightIdx)
    LeftShift = np.roll(lightIdx, -1)  # 向左移动一位
    # print("LeftShift:", LeftShift)
    differList = lightIdx - LeftShift
    # print("differList:", differList)
    temp = differList.tolist()
    col_gapLeft = lightIdx[temp.index(min(differList))]

    # 再将间断点右边的 光条排除
    laser1_ = laser1_threshold[:, : col_gapLeft + 1]  # 刚好截止到 左间断点

    lightList_part = np.sum(laser1_, axis=1)  # 1 行相加, 变成一列
    ret = np.where(lightList_part > 0)  # 返回（索引+ 数据类型 ）
    lightIdx = ret[0]
    rowIdx = lightIdx[-1]

    # 计算这一行的特征点坐标
    _, width = laser1_.shape
    n, sumCol = 0, 0
    for col in range(width):  # 遍历这一行的每一列
        if laser1_[rowIdx, col] != 0:
            n += 1
            sumCol += (col + 1)  # 重心在第几列
    indexCol = round(sumCol / n) - 1  # 化为坐标索引

    P1 = [indexCol, rowIdx]  # 行索引，列索引
    print("laser1_的焊缝特征点P1：", P1)
    # 基于P1，取出P1右边部分的 一块ROI
    roiRight_col = col_gapLeft + 80  # col_gapLeft :间隙左边点的列索引
    if roiRight_col > laser1_threshold.shape[1] - 1:
        roiRight_col = laser1_threshold.shape[1] - 1
    ROI = laser1_threshold[rowIdx - 20:rowIdx + 80, col_gapLeft + 1:roiRight_col + 1]  # col_gapLeft + 1：刚好没有光点
    cv.rectangle(L1img, (col_gapLeft + 1, rowIdx - 20), (roiRight_col, rowIdx + 79), (0, 255, 0), 1)  # 列，行

    # 水平COG提取ROI内的光条中心点坐标
    # ROI左上角 起点坐标 列：col_gapLeft + 1 行：rowIdx - 20
    startP = (rowIdx - 20, col_gapLeft + 1)
    # Py = startP[0], Px = startP[1]  # 行，列
    cp_row, cp_col = cog(ROI, startP, 1)  # 1 : 水平cog

    # 将sp1 和 P1连接，得到laser1 的左边部分的直线方程L1 斜率
    x_P1, y_P1 = P1
    x_sp1, y_sp1 = start_point1  # （列，行）
    kl = (y_P1 - y_sp1) / (x_P1 - x_sp1)  # 斜率

    # 然后直线的斜率，结合EP1点，得到laser1 的右边部分的直线方程L2
    (x0, y0) = EP1
    A = kl
    B = -1
    C = y0 - kl * x0
    # 计算点到L2距离，分离出坡口上的中心点，再拟合2段点集，求出交点P1_1
    LineRow, LineCol = [], []
    grooveRow, grooveCol = [], []
    countLine, countGroove = 0, 0
    for i in range(len(cp_row)):
        d = math.fabs(A * cp_col[i] + B * cp_row[i] + C) / math.sqrt(A ** 2 + B ** 2)
        if d <= 3:
            LineRow += [cp_row[i]]
            LineCol += [cp_col[i]]
            countLine += 1
        if d > 4:
            grooveRow += [cp_row[i]]
            grooveCol += [cp_col[i]]
            countGroove += 1
    print("坡口%d个点" % countGroove)
    print("非坡口光条%d个点" % countLine)

    Al, Bl, Cl = fitLine((LineRow, LineCol))  # pointSet_row, pointSet_col = pointSet
    Agl, Bgl, Cgl = fitLine((grooveRow, grooveCol))
    # 计算交点
    A1 = np.array([[Al, Bl], [Agl, Bgl]])  # 输出系数矩阵A
    b1 = np.array([-Cl, -Cgl])  # 值
    fp_x1, fp_y1 = solve(A1, b1)
    fp_y1, fp_x1 = round(fp_y1), round(fp_x1)

    P1_w = [fp_x1, fp_y1]  # 行索引，列索引

    cv.circle(L1img, P1, 1, (255, 0, 255), -1)
    cv.circle(L1img, P1, 5, (255, 0, 255), 1)

    cv.circle(L1img, P1_w, 1, (255, 0, 255), -1)
    cv.circle(L1img, P1_w, 5, (255, 0, 255), 1)

    # cv.imshow('P1_L1img', L1img)
    cv.imwrite(r"F:\MVS_Data\P1_L1img.png", L1img)

    return P1, P1_w


def laser2_Extract(L2img, start_point2, EP2):
    print("********************** laser2_ 特征提取 **********************")
    Red = L2img[:, :, 2]  # B0 G1 R2
    _, threshold = cv.threshold(Red, 190, 255, cv.THRESH_BINARY)  # 转为二值图
    laser2_threshold = cv.medianBlur(threshold, 3)
    laser2_threshold_copy = laser2_threshold

    # cv.imshow('laser2_threshold', laser2_threshold)
    cv.imwrite(r"F:\MVS_Data\laser2_threshold.jpg", laser2_threshold)

    hight, width = laser2_threshold.shape

    # 将下面的重叠光条去掉
    for col in range(width):
        n = 0
        for row in range(hight - 1):
            diff = int(laser2_threshold[row + 1, col]) - int(laser2_threshold[row, col])  # 下一行像素值减去上一行像素值
            if diff > 0:
                n += 1
            if n > 1 and laser2_threshold[row, col] != 0:  # 将上面的重叠光条去掉
                laser2_threshold[row, col] = 0
    # cv.imshow('remove', laser2_threshold)
    cv.waitKey(0)

    # 求laser1光条的下间断点行坐标
    lightList = np.sum(laser2_threshold, axis=1)  # 1 行相加, 变成一列
    ret = np.where(lightList > 0)  # 返回（索引+ 数据类型 ）
    lightIdx = ret[0]
    # print("lightIdx:", lightIdx)
    LeftShift = np.roll(lightIdx, -1)  # 向下移动一位
    # print("LeftShift:", LeftShift)
    differList = lightIdx - LeftShift
    # print("differList:", differList)
    temp = differList.tolist()
    row_gapUp = lightIdx[temp.index(min(differList))]

    # 计算这一行的特征点坐标
    n, sumRow = 0, 0
    for col in range(width):  # 遍历这一行的每一列
        if laser2_threshold[row_gapUp, col] != 0:
            n += 1
            sumRow += (col + 1)  # 重心在第几列
    indexCol = round(sumRow / n) - 1  # 化为坐标索引

    P2 = [indexCol, row_gapUp]
    print("laser2_的焊缝特征点P2：", P2)
    # 基于P2，取出P2下边的 一块ROI
    roiLeft_col = indexCol - 80
    if roiLeft_col < 0:
        roiLeft_col = 0
    ROI = laser2_threshold_copy[row_gapUp + 1:row_gapUp + 100, roiLeft_col:indexCol + 30]
    cv.rectangle(L2img, (roiLeft_col, row_gapUp + 1), (indexCol + 29, row_gapUp + 99), (0, 255, 0), 1)  # 列，行

    # 垂直COG提取ROI内的光条中心点坐标
    # ROI左上角 起点坐标 列：roiLeft_col 行：row_gapUp + 1
    startP = (row_gapUp + 1, roiLeft_col)
    # Py = startP[0], Px = startP[1]  # 行，列
    cp_row, cp_col = cog(ROI, startP, 0)  # 0 : 垂直cog

    # 将EP2 和 P2连接，得到laser2 的右边部分的直线方程L1 斜率
    x_P2, y_P2 = P2
    x_EP2, y_EP2 = EP2  # （列，行）
    kl = (y_P2 - y_EP2) / (x_P2 - x_EP2)  # 斜率

    # 然后直线的斜率，结合sp2点，得到laser1 的左边部分的直线方程L2
    (x0, y0) = start_point2
    A = kl
    B = -1
    C = y0 - kl * x0
    # 计算点到L2距离，分离出坡口上的中心点，再拟合2段点集，求出交点P1_1
    LineRow, LineCol = [], []
    grooveRow, grooveCol = [], []
    countLine, countGroove = 0, 0
    distance = []
    for i in range(len(cp_row)):
        d = math.fabs(A * cp_col[i] + B * cp_row[i] + C) / math.sqrt(A ** 2 + B ** 2)
        distance += [d]
        if d <= 3:
            LineRow += [cp_row[i]]
            LineCol += [cp_col[i]]
            countLine += 1
        if d > 4:
            grooveRow += [cp_row[i]]
            grooveCol += [cp_col[i]]
            countGroove += 1
    print("坡口%d个点" % countGroove)
    print("非坡口光条%d个点" % countLine)
    print("distance:", distance)

    Al, Bl, Cl = fitLine((LineRow, LineCol))  # pointSet_row, pointSet_col = pointSet
    Agl, Bgl, Cgl = fitLine((grooveRow, grooveCol))
    # 计算交点
    A1 = np.array([[Al, Bl], [Agl, Bgl]])  # 输出系数矩阵A
    b1 = np.array([-Cl, -Cgl])  # 值
    fp_x2, fp_y2 = solve(A1, b1)
    fp_y2, fp_x2 = round(fp_y2), round(fp_x2)

    P2_w = [fp_x2, fp_y2]  #

    cv.circle(L2img, P2, 1, (255, 0, 255), -1)
    cv.circle(L2img, P2, 5, (255, 0, 255), 1)

    cv.circle(L2img, P2_w, 1, (255, 0, 255), -1)
    cv.circle(L2img, P2_w, 5, (255, 0, 255), 1)

    # cv.imshow('P2_L2img', L2img)
    cv.imwrite(r"F:\MVS_Data\P2_L2img.png", L2img)

    return P2, P2_w


def cog(roi, startP, mode):
    h, w = roi.shape
    Py, Px = startP  # 行，列
    cp_row, cp_col = [], []
    if mode == 0:  # 垂直COG
        for col in range(w):
            n = 0
            sum_RowIdx = 0
            for row in range(h):
                if roi[row, col] != 0:
                    n += 1
                    sum_RowIdx += row + 1
            if n:
                indexRow = round(sum_RowIdx / n) - 1
                cp_col += [col + Px]
                cp_row += [indexRow + Py]
    elif mode == 1:  # 水平COG
        for row in range(h):
            n = 0
            sum_ColIdx = 0
            for col in range(w):
                if roi[row, col] != 0:
                    n += 1
                    sum_ColIdx += col + 1
            if n:
                indexCol = round(sum_ColIdx / n) - 1
                cp_col += [indexCol + Px]
                cp_row += [row + Py]
    return cp_row, cp_col
